import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import label
from skimage import exposure


def extract_stripes_from_fft(image_path, min_peak_distance=10, min_peak_value=0.1):
    """
    对周期性条纹图片进行傅里叶变换分析，提取对称亮点并重建条纹图像

    参数:
    image_path (str): 图像文件路径
    min_peak_distance (int): 检测亮点的最小距离
    min_peak_value (float): 检测亮点的最小相对强度 (0-1)

    返回:
    list: 重建的条纹图像列表
    dict: 包含FFT结果的字典
    """
    # 1. 读取图像并预处理
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise ValueError(f"无法加载图像: {image_path}")

    # 对比度增强
    p2, p98 = np.percentile(img, (2, 98))
    img = exposure.rescale_intensity(img, in_range=(p2, p98))

    # 2. 傅里叶变换
    f = np.fft.fft2(img)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = np.log(np.abs(fshift) + 1)

    # 3. 检测对称亮点
    center = np.array(magnitude_spectrum.shape) // 2
    peaks = find_symmetric_peaks(magnitude_spectrum, center, min_peak_distance, min_peak_value)

    # 4. 对每对亮点进行逆变换
    reconstructed_images = []
    for i, (peak1, peak2) in enumerate(peaks):
        # 创建只包含当前亮点对的掩码
        mask = np.zeros_like(fshift, dtype=bool)
        mask = add_peak_mask(mask, peak1, magnitude_spectrum.shape)
        mask = add_peak_mask(mask, peak2, magnitude_spectrum.shape)

        # 应用掩码进行逆变换
        filtered_fft = fshift * mask
        filtered_fft_ishift = np.fft.ifftshift(filtered_fft)
        reconstructed = np.abs(np.fft.ifft2(filtered_fft_ishift))

        # 归一化并保存
        reconstructed = exposure.rescale_intensity(reconstructed, out_range=(0, 255))
        reconstructed_images.append(reconstructed.astype(np.uint8))

    # 可视化结果
    visualize_results(img, magnitude_spectrum, peaks, reconstructed_images)

    return reconstructed_images, {
        'original': img,
        'magnitude_spectrum': magnitude_spectrum,
        'peaks': peaks
    }


def find_symmetric_peaks(magnitude_spectrum, center, min_distance=10, min_value=0.1):
    """
    在傅里叶频谱中检测对称亮点对
    """
    # 归一化频谱 (0-1)
    norm_spec = magnitude_spectrum - np.min(magnitude_spectrum)
    norm_spec = norm_spec / np.max(norm_spec)

    # 创建中心掩码 (排除中心区域)
    center_mask = np.zeros_like(norm_spec, dtype=bool)
    center_radius = min(center) // 4
    y, x = np.ogrid[:magnitude_spectrum.shape[0], :magnitude_spectrum.shape[1]]
    center_mask = (x - center[1]) ** 2 + (y - center[0]) ** 2 <= center_radius ** 2

    # 检测亮点
    threshold = min_value * np.max(norm_spec)
    peak_mask = (norm_spec > threshold) & ~center_mask

    # 标记连通区域
    labeled, num_features = label(peak_mask)
    peak_props = []

    # 计算每个亮点的重心和强度
    for i in range(1, num_features + 1):
        region_mask = labeled == i
        y_indices, x_indices = np.where(region_mask)
        intensity = np.sum(norm_spec[region_mask])
        centroid = (np.mean(y_indices), np.mean(x_indices))
        peak_props.append((centroid, intensity))

    # 按强度排序
    peak_props.sort(key=lambda x: x[1], reverse=True)

    # 寻找对称点对
    peaks = []
    used = set()

    for i, (pos1, intensity1) in enumerate(peak_props):
        if i in used:
            continue

        # 寻找对称点
        symmetric_pos = 2 * center - np.array(pos1)
        symmetric_pos = tuple(np.round(symmetric_pos).astype(int))

        # 检查对称点是否在有效范围内
        if (symmetric_pos[0] < 0 or symmetric_pos[0] >= magnitude_spectrum.shape[0] or
                symmetric_pos[1] < 0 or symmetric_pos[1] >= magnitude_spectrum.shape[1]):
            continue

        # 寻找最近的亮点作为对称点
        best_match = None
        best_dist = float('inf')

        for j, (pos2, intensity2) in enumerate(peak_props):
            if j == i or j in used:
                continue

            dist = np.linalg.norm(np.array(pos2) - symmetric_pos)
            if dist < min_distance and dist < best_dist:
                best_match = j
                best_dist = dist

        # 如果找到匹配，添加点对
        if best_match is not None:
            peaks.append((pos1, peak_props[best_match][0]))
            used.add(i)
            used.add(best_match)

    return peaks


def add_peak_mask(mask, peak, shape, radius=3):
    """
    在掩码中添加峰值区域
    """
    y, x = peak
    y_min = max(0, int(y) - radius)
    y_max = min(shape[0], int(y) + radius + 1)
    x_min = max(0, int(x) - radius)
    x_max = min(shape[1], int(x) + radius + 1)

    mask[y_min:y_max, x_min:x_max] = True
    return mask


def visualize_results(original, magnitude_spectrum, peaks, reconstructed_images):
    """
    可视化处理结果
    """
    plt.figure(figsize=(15, 10))

    # 原始图像
    plt.subplot(231)
    plt.imshow(original, cmap='gray')
    plt.title('原始图像')
    plt.axis('off')

    # 傅里叶频谱
    plt.subplot(232)
    plt.imshow(magnitude_spectrum, cmap='viridis')
    plt.title('傅里叶频谱')
    plt.axis('off')

    # 标记亮点的频谱
    plt.subplot(233)
    plt.imshow(magnitude_spectrum, cmap='viridis')
    center = np.array(magnitude_spectrum.shape) // 2
    plt.scatter(center[1], center[0], c='red', s=50, marker='x')

    for peak1, peak2 in peaks:
        plt.scatter(peak1[1], peak1[0], c='red', s=50)
        plt.scatter(peak2[1], peak2[0], c='red', s=50)
        plt.plot([peak1[1], peak2[1]], [peak1[0], peak2[0]], 'r-')

    plt.title('检测到的对称亮点')
    plt.axis('off')

    # 重建的条纹图像
    for i, img in enumerate(reconstructed_images[:3]):
        plt.subplot(2, 3, i + 4)
        plt.imshow(img, cmap='gray')
        plt.title(f'重建条纹 {i + 1}')
        plt.axis('off')

    plt.tight_layout()
    plt.show()


# 使用示例
if __name__ == "__main__":
    image_path = "../img/pic1.jpg"  # 替换为你的条纹图像路径
    reconstructed_images, fft_data = extract_stripes_from_fft(image_path)

    # 显示所有重建的条纹图像
    plt.figure(figsize=(15, 5))
    for i, img in enumerate(reconstructed_images):
        plt.subplot(1, len(reconstructed_images), i + 1)
        plt.imshow(img, cmap='gray')
        plt.title(f'条纹分量 {i + 1}')
        plt.axis('off')
    plt.tight_layout()
    plt.show()